#include <iostream>
#include <opencv2/opencv.hpp>

// Include YOLO11 detector
#include "det/YOLO11_RDK.hpp"

using namespace yolos_edgeplatform;

int main(int argc, char** argv) {
    if (argc < 3) {
        std::cout << "Usage: " << argv[0] << " <model_path> <image_path> [labels_path]" << std::endl;
        std::cout << "Example: " << argv[0] << " model.bin test.jpg coco.names" << std::endl;
        return -1;
    }

    std::string modelPath = argv[1];
    std::string imagePath = argv[2];
    std::string labelsPath = argc > 3 ? argv[3] : "coco.names";

    try {
        // Initialize detector
        std::cout << "=== Initializing YOLO11 Detector ===" << std::endl;
        YOLO11Detector detector(modelPath, labelsPath);

        // Load image
        cv::Mat image = cv::imread(imagePath);
        if (image.empty()) {
            std::cerr << "[ERROR] Failed to load image: " << imagePath << std::endl;
            return -1;
        }

        std::cout << "\n=== Running Detection ===" << std::endl;
        std::cout << "Image size: " << image.cols << "x" << image.rows << std::endl;

        // Perform detection
        auto detections = detector.detect(image, 0.25f, 0.45f);

        std::cout << "\n=== Detection Results ===" << std::endl;
        std::cout << "Found " << detections.size() << " objects" << std::endl;

        // Draw results
        const auto &classNames = detector.getClassNames();
        for (const auto &det : detections) {
            // Draw bounding box
            cv::rectangle(image,
                         cv::Point(det.box.x, det.box.y),
                         cv::Point(det.box.x + det.box.width, det.box.y + det.box.height),
                         cv::Scalar(0, 255, 0), 2);

            // Draw label
            std::string label = classNames[det.classId] + ": " +
                               std::to_string(static_cast<int>(det.conf * 100)) + "%";
            cv::putText(image, label,
                       cv::Point(det.box.x, det.box.y - 5),
                       cv::FONT_HERSHEY_SIMPLEX, 0.5,
                       cv::Scalar(0, 0, 255), 1, cv::LINE_AA);

            // Print detection info
            std::cout << "  - " << label << " at ("
                     << static_cast<int>(det.box.x) << ", "
                     << static_cast<int>(det.box.y) << ", "
                     << static_cast<int>(det.box.width) << ", "
                     << static_cast<int>(det.box.height) << ")" << std::endl;
        }

        // Save result
        std::string outputPath = "result.jpg";
        cv::imwrite(outputPath, image);
        std::cout << "\n[INFO] Result saved to: " << outputPath << std::endl;

    } catch (const std::exception &e) {
        std::cerr << "[ERROR] " << e.what() << std::endl;
        return -1;
    }

    return 0;
}
